import wmi
import winreg
import hashlib
import pythoncom
from hardware.edid import edid
import logging


class HardwareUtil:
    def __init__(self):
        pythoncom.CoInitialize()
        self.connect = wmi.WMI()

    def __del__(self):
        try:
            pythoncom.CoUninitialize()
        except Exception as e:
            print(e)
            logging.error(e)

    # 获取系统信息
    def system(self):
        _keys = ["Manufacturer", "SystemFamily", "SystemType", "Model", "Name", "NameFormat", "NumberOfProcessors", "NumberOfLogicalProcessors"]
        _r = formatData(self.connect.Win32_ComputerSystem(_keys)[0])
        _r['TYPE'] = 'SYSTEM'
        if 'Manufacturer' in _r and 'SystemFamily' in _r:
            _r['name'] = "{} {}".format(_r['Manufacturer'], _r['SystemFamily'])
        elif 'SystemFamily' in _r:
            _r['name'] = _r['SystemFamily']
        elif 'SystemType' in _r:
            _r['name'] = _r['SystemType']
        else:
            _r['name'] = _r['Name']
        return _r

    # 获取主板信息
    def baseboard(self):
        _keys = ["Manufacturer", "Model", "Name", "Product", "SerialNumber"]
        _r = formatData(self.connect.Win32_BaseBoard(_keys)[0])
        _r['TYPE'] = 'BASEBOARD'
        if 'Manufacturer' in _r and 'Product' in _r:
            _r['name'] = "{} {}".format(_r['Manufacturer'], _r['Product'])
        elif 'Product' in _r:
            _r['name'] = _r['Product']
        else:
            _r['name'] = "Base Board"
        return _r

    # 获取主板信息
    def bios(self):
        _keys = ["Manufacturer", "SerialNumber", "SMBIOSBIOSVersion", "Version", "Name"]
        _r = formatData(self.connect.Win32_BIOS(_keys)[0])
        _r['TYPE'] = 'BIOS'
        if 'SMBIOSBIOSVersion' in _r:
            _r['name'] = str(_r['SMBIOSBIOSVersion'])
        elif 'Manufacturer' in _r:
            _r['name'] = str(_r['Manufacturer'])
        else:
            _r['name'] = "Default BIOS"
        return _r

    # 获取CPU列表
    def cpu(self):
        _keys = ["Caption", "CurrentClockSpeed", "CurrentVoltage", "DataWidth", "Description", "DeviceID", "ExtClock", "ProcessorType",
                 "Family", "L2CacheSize", "L3CacheSize", "L3CacheSpeed", "Level", "Manufacturer", "MaxClockSpeed", "Name",
                 "NumberOfCores", "NumberOfEnabledCore", "NumberOfLogicalProcessors", "ProcessorId", "SerialNumber",
                 "SocketDesignation", "SystemName", "ThreadCount", "Version"]
        _rs = []
        for _d in self.connect.Win32_Processor(_keys):
            _r = formatData(_d)
            _r['TYPE'] = 'CPU'
            if 'Name' in _r:
                _r['name'] = _r['Name']
            elif 'Manufacturer' in _r:
                _r['name'] = _r['Manufacturer']
            else:
                _r['name'] = "Default CPU"
            _rs.append(_r)
        return _rs

    # 获取内存列表
    def memory(self):
        _keys = ["BankLabel", "Capacity", "ConfiguredClockSpeed", "ConfiguredVoltage", "DataWidth", "Description", "TotalWidth", "TypeDetail",
                 "PartNumber", "Speed", "Manufacturer", "MaxVoltage", "MinVoltage", "Name", "PartNumber", "SerialNumber", "SMBIOSMemoryType", "Speed"]
        _rs = []
        for _d in self.connect.Win32_PhysicalMemory(_keys):
            _r = formatData(_d)
            _r['TYPE'] = 'MEMORY'
            if 'Manufacturer' in _r and 'SMBIOSMemoryType' in _r and 'Speed' in _r and 'Capacity' in _r:
                _r['name'] = "{} {} {}MHz ({})".format(_r['Manufacturer'], self.parseMemoryType(_r['SMBIOSMemoryType']), _r['Speed'], self.parseBits(_r['Capacity']))
            elif 'Manufacturer' in _r and 'SMBIOSMemoryType' in _r and 'Capacity' in _r:
                _r['name'] = "{} {} ({})".format(_r['Manufacturer'], self.parseMemoryType(_r['SMBIOSMemoryType']), self.parseBits(_r['Capacity']))
            elif 'Manufacturer' in _r and 'Capacity' in _r:
                _r['name'] = "{} {}".format(_r['Manufacturer'], self.parseBits(_r['Capacity']))
            elif 'Manufacturer' in _r:
                _r['name'] = _r['Manufacturer']
            else:
                _r['name'] = "Default Memory"
            _rs.append(_r)
        return _rs

    # 获取硬盘
    def disk(self):
        _keys = ["Caption", "Description", "FirmwareRevision", "InterfaceType", "Manufacturer", "MediaType", "Model", "Name", "Size", "SerialNumber"]
        _rs = []
        for _d in self.connect.Win32_DiskDrive(_keys):
            mainDisk = False
            for _partition in _d.associators("Win32_DiskDriveToDiskPartition"):
                if _partition.BootPartition:
                    mainDisk = True
            _r = formatData(_d)
            _r['BootPartition'] = mainDisk
            _r['TYPE'] = 'DISK'
            if 'Model' in _r and 'Size' in _r:
                _r['name'] = "{} ({})".format(_r['Model'], self.parseBits(_r['Size']))
            else:
                _r['name'] = _r['Model']
            _rs.append(_r)
        return _rs

    # 获取显卡信息
    def gpu(self):
        _keys = ["AdapterCompatibility", "AdapterDACType", "AdapterRAM", "Caption", "CurrentBitsPerPixel", "CurrentHorizontalResolution", "VideoProcessor",
                 "CurrentNumberOfColors", "CurrentRefreshRate", "Description", "MaxRefreshRate", "MinRefreshRate", "Name", "VideoArchitecture", "VideoMemoryType", "VideoModeDescription"]
        _rs = []
        for _d in self.connect.Win32_VideoController(_keys, Availability=3):
            if _d.VideoProcessor is None:
                continue
            _r = formatData(_d)
            _r['TYPE'] = 'GraphicsCard'
            if 'Name' in _r and 'AdapterRAM' in _r:
                _r['name'] = "{} ({})".format(_r['Name'], self.parseBits(_r['AdapterRAM']))
            else:
                _r['name'] = _r['Name']
            _rs.append(_r)
        return _rs

    # 获取显示器信息
    def display(self):
        displays = []
        for d in self.connect.Win32_DesktopMonitor():
            try:
                device_opend = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, "SYSTEM\\ControlSet001\\Enum\\" + d.PNPDeviceID + "\\Device Parameters")
                edid_bytes = winreg.QueryValueEx(device_opend, "EDID")
                display = edid(edid_bytes[0])
                _r = display.__dict__
                _r['TYPE'] = 'DISPLAY'
                if 'name' in _r and 'serial' in _r:
                    _r['name'] = "{} ({})".format(_r['name'], _r['serial'])
                elif 'name' in _r:
                    _r['name'] = str(_r['name'])
                elif 'serial' in _r:
                    _r['name'] = str(_r['serial'])
                else:
                    _r['name'] = 'Default Display'
                displays.append(_r)
            except Exception as e1:
                displays.append({'TYPE': 'DISPLAY', 'name': d.Name})
                logging.error(e1)
                continue
        return displays

    # 获取声音设备信息
    def sound(self):
        _keys = ["Caption", "Name", "Manufacturer", "Description", "ProductName"]
        sounds = []
        for _d in self.connect.Win32_SoundDevice(_keys):
            if _d.Manufacturer != "(Generic USB Audio)":
                _r = formatData(_d)
                _r['TYPE'] = 'SOUND'
                _r['name'] = _r['Name']
                sounds.append(_r)
        return sounds

    # 获取键盘信息
    def keyboard(self):
        _keys = ["Caption", "Name", "Manufacturer", "Description"]
        keyboard = self.connect.Win32_PNPEntity(_keys, PNPClass="Keyboard")
        if len(keyboard) > 0:
            _r = formatData(keyboard[0])
            _r['TYPE'] = 'KEYBOARD'
            _r['name'] = _r['Name']
            return _r
        return None

    # 获取鼠标信息
    def mouse(self):
        _keys = ["Caption", "Name", "Manufacturer", "Description"]
        mouse = self.connect.Win32_PNPEntity(_keys, PNPClass="Mouse")
        if len(mouse) > 0:
            _r = formatData(mouse[0])
            _r['TYPE'] = 'MOUSE'
            _r['name'] = _r['Name']
            return _r
        return None

    # 获取网卡信息
    def network(self):
        _keys = ["Caption", "Name", "Manufacturer", "Description"]
        networks = []
        for _d in self.connect.Win32_PNPEntity(_keys, PNPClass="Net"):
            if _d.Manufacturer != 'Microsoft':
                _r = formatData(_d)
                _r['TYPE'] = 'NETWORK'
                _r['name'] = _r['Name']
                networks.append(_r)
        return networks

    # 获取播放设备
    def media(self):
        _keys = ["Caption", "Name", "Manufacturer", "Description"]
        _rs = []
        _ds = self.connect.Win32_PNPEntity(_keys, PNPClass="MEDIA", Service="usbaudio")
        for _d in _ds:
            _r = formatData(_d)
            _r['TYPE'] = 'MEDIA'
            _r['name'] = _r['Name']
            _rs.append(_r)
        return _rs

    # 获取Windows激活码
    def productkey(self):
        _r = {'TYPE': 'PRODUCTKEY'}
        try:
            platform = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE,
                                      "SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\SoftwareProtectionPlatform")
            productKey = winreg.QueryValueEx(platform, "BackupProductKeyDefault")
            _r['name'] = productKey[0]
        except Exception as e1:
            logging.error(e1)
            _r['name'] = 'NONE'
        _d = self.connect.SoftwareLicensingService(["OA3xOriginalProductKey"])[0]
        _d = formatData(_d)
        if _d['OA3xOriginalProductKey'] is not None and _d['OA3xOriginalProductKey'] != '':
            _r['name'] = _d['OA3xOriginalProductKey']
        return _r

    # 格式化内存版本号
    def parseMemoryType(self, type):
        types = {
            0x01: "Other",
            0x03: "DRAM",
            0x04: "EDRAM",
            0x05: "VRAM",
            0x06: "SRAM",
            0x07: "RAM",
            0x08: "ROM",
            0x09: "FLASH",
            0x0A: "EEPROM",
            0x0B: "FEPROM",
            0x0C: "EPROM",
            0x0D: "CDRAM",
            0x0E: "3DRAM",
            0x0F: "SDRAM",
            0x10: "SGRAM",
            0x11: "RDRAM",
            0x12: "DDR",
            0x13: "DDR2",
            0x14: "DDR2 FB-DIMM",
            0x18: "DDR3",
            0x19: "FBD2",
            0x1A: "DDR4",
            0x1B: "LPDDR",
            0x1C: "LPDDR2",
            0x1D: "LPDDR3",
            0x1E: "LPDDR4",
            0x1F: "Logical non-volatile device",
        }
        return types.get(int(type), None)

    def parseBits(self, capacity):
        if capacity is None:
            return "Zero bit"
        capacity = int(capacity)
        exts = {
            0: "bit",
            1: "KB",
            2: "MB",
            3: "GB",
            4: "TB"
        }
        i = 0
        while capacity >= 1024:
            i = i + 1
            capacity = capacity / 1024
        return str(round(capacity)) + "" + exts.get(i, "GB")


def formatData(_d):
    _r = {}
    for _k in _d.properties:
        _v = _d.__getattr__(_k)
        if _v is None:
            continue
        if _k == 'PartNumber':
            _r['PartModel'] = hashlib.sha1(_v.encode('utf8')).hexdigest()
        elif _k == 'ProcessorId':
            _r['ProcessorModel'] = hashlib.sha1(_v.encode('utf8')).hexdigest()
        elif _k == 'SerialNumber':
            _r['ProductModel'] = hashlib.sha1(_v.encode('utf8')).hexdigest()
        else:
            _r[_k] = str(_v)
    return _r


if __name__ == "__main__":
    print(HardwareUtil().display())
